function [C,D,resid,invtausq_d,llik_trace] = binary_matrix_factorization(B,K,BURNIN,NOPT)
%
% [C,D,resid,invtausq_d,llik] = binary_matrix_factorization(B,K,BURNIN,NOPT)
%
% factorize B ~ sigmoid(C * D')
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    if any(B(:) > 1) || any(B(:) < 0),
        error('B must be binary matrix in [0,1]');
    end

    assert(K > 1, 'at least two factors (including one constant factor)');

    TOL = 1e-4;
    PTRUNC = 1e-4;
    
    missing = isnan(B);
    [m, n] = size(B);

    B(missing) = 0;
    B = single(B);

    % randomly initialize C
    C = 0.1*randn(m,K,'single');
    D = zeros(n,K,'single');

    % last column is just constant
    D(:,K) = 1;

    invtausq_d = zeros(1,K,'single');
    invtausq_c = 0.01*K*ones(m,1,'single');

    % ================================================================
    llik_trace = NaN(BURNIN+NOPT,1,'single');

    for iter = 1:(BURNIN+NOPT),

        % 1. quadratic approximation
        F = C*D';
        P = 1./(1+exp(-F));
        
        P(P < PTRUNC) = PTRUNC;
        P(P > 1 - PTRUNC) = 1 - PTRUNC;

        W = P.*(1-P);
        X = F + (B - P) ./ W;

        W(missing) = 0;

        % 2. estimation D
        %
        %          sum_i W(i,j) * [ X(i,j) - sum_k' C(i,k') * D(j,k') + C(i,k) * D(j,k) ] * C(i,k)
        % D(j,k) = -------------------------------------------------------------------------------
        %          sum_i W(i,j) * C(i,k)^2
        % 

        WtC2 = W'*(C.^2);
        R = X - F;

        % update each D(:,k) exept the last one
        for k = 1:(K-1),
            R_exc = R + bsxfun(@times, C(:,k), D(:,k).');
            denom_d = WtC2(:,k) + invtausq_d(k);
            num_d = (W.*R_exc)' * C(:,k);

            dd = num_d ./ denom_d;
            D(:,k) = dd;

            R = R_exc - bsxfun(@times, C(:,k), D(:,k).');
        end
        F = C*D';

        % 3. estimate C
        %
        %          sum_j W(i,j) * [ X(i,j) - sum_k' C(i,k') * D(j,k') + C(i,k) * D(j,k) ] * D(j,k)
        % C(i,k) = -------------------------------------------------------------------------------
        %          sum_j W(i,j) * D(j,k)^2
        %

        WD2 = W * (D.^2);
        for k = 1:K,
            R_exc = R + bsxfun(@times, C(:,k), D(:,k).');
            denom_c = WD2(:,k) + invtausq_c;
            num_c = (W.*R_exc) * D(:,k);

            cc = num_c ./ denom_c;
            C(:,k) = cc;

            R = R_exc - bsxfun(@times, C(:,k), D(:,k).');
        end
        F = C*D';

        llik = sum(sum(B.*F - log(1+exp(F))));
        llik_trace(iter) = llik;

        if iter > BURNIN,
            % measure variance of D
            d_var = 1./bsxfun(@plus, W'*(C.^2), invtausq_d);

            % estimate lambda
            % lambda = n * K / sum_k sqrt{d(:,k)'*d(:,k)}
            dsq_mat = D.^2 + d_var;
            dsq_mat = arrayfun(@(x) max(PTRUNC, x), dsq_mat);

            lambda = n * (K-1) ./ sum(sqrt(sum(dsq_mat(:,1:(K-1)))));
            invtausq_d = lambda ./ sqrt(sum(dsq_mat));
        end

        if iter > BURNIN + 10,
            llik_prev = mean(llik_trace((iter-10):(iter-6)));
            llik_curr = mean(llik_trace((iter-5):(iter-1)));

            if abs(llik_prev - llik_curr) < TOL,
                llik_trace = llik_trace(1:iter);
                fprintf(1,'Converged\n');
                break;
            end
        end
        
        fprintf(1,'Iter = %03d, LLIK = %.4e\n', iter, llik);
    end

    resid = 1./(1+exp(-(X-F)));
end
